import numpy as np
from scipy.interpolate import griddata

import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D


file_name = "track_pts_02_cleaned"
data = np.loadtxt(file_name + ".txt")
x = data[:,0]
y = data[:,1]
imgsize = (640,480)
bin_size = 20

r = [[-.5*bin_size, imgsize[0]+.5*bin_size], [-.5*bin_size, imgsize[1]+.5*bin_size]]

heatmap, xedges, yedges = np.histogram2d(x, y, bins=(1+imgsize[0]/bin_size,1+imgsize[1]/bin_size), range=r)

#offset gridpoints to the bin ceters
xedges[:-1] +=  0.5*bin_size
yedges[:-1] +=  0.5*bin_size

# correlate data point with heatmap value
xmesh, ymesh = np.meshgrid(xedges[:-1], yedges[:-1])
grid = np.hstack((xmesh.flatten()[:,None],ymesh.flatten()[:,None]))
heatmap = heatmap.transpose()
heatmap_max = np.amax(heatmap)

# grid the data.
xi = np.arange(0,xedges[-1],1)
yi = np.arange(0,yedges[-1],1)

zi = griddata(grid, heatmap.flatten(), (xi[None,:], yi[:,None]), method='cubic')

# contour the gridded data, plotting dots at the randomly spaced data points.
fig = plt.figure()
ax = fig.gca(projection='3d')

X, Y = np.meshgrid(xi, yi)
ax.contour(X, Y, zi, range(1,int(heatmap_max),10), linewidths=0.5,colors='r')
ax.plot_surface(X, Y, zi, rstride=10, cstride=10, cmap=plt.cm.jet, vmin=1, vmax=int(heatmap_max), alpha=0.5)
#ax.plot_wireframe(X, Y, zi, rstride=10, cstride=10)
ax.set_zlim3d(1, int(heatmap_max))
plt.savefig(file_name + '3d.svg', format='svg')
plt.show()
'''
plt.scatter(x,y, c=(0,1,.25), s=2, marker='+', linestyle='solid')
plt.ylim(0, imgsize[1])
plt.xlim(0, imgsize[0])
plt.plot(x,y, c=(0.5,0,1.0), linestyle='solid')
#plt.scatter(grid[:,0], grid[:,1], c=(1,1,1), s=2)
#plt.colorbar() # draw colorbar
plt.savefig(file_name + '3d.svg', format='svg')
plt.show()
'''